package com.lql;


import java.io.*;
import java.util.*;

/**
 * 实现赫夫曼编码
 */
public class HuffmanTreeCode {
    public static void main(String[] args) {
//        String strs = "i like like like java do you like a java hahazxc";
//        byte[] bytes = huffmanZip(strs);
//        System.out.println(Arrays.toString(bytes));
//        System.out.println(byteToBinaryString(bytes));
//        System.out.println(new String(unZipHuffmanCode(huffmanCodeMap, byteToBinaryString(bytes))));
//        fileHuffmanZip("C:\\Users\\lql\\Desktop\\test.txt", "C:\\Users\\lql\\Desktop\\test1.zip");
//        unzipFile("C:\\Users\\lql\\Desktop\\test1.zip", "C:\\Users\\lql\\Desktop\\test1.txt");
    }

    /**
     * 解压文件
     *
     * @param unzipFilePath    待解压文件的全路径
     * @param desUnzipFilePath 解压到的目的全路径
     */
    public static void unzipFile(String unzipFilePath, String desUnzipFilePath) {
        FileInputStream fileInputStream = null;
        ObjectInputStream objectInputStream = null;
        FileOutputStream fileOutputStream = null;
        try {
            fileInputStream = new FileInputStream(unzipFilePath);
            objectInputStream = new ObjectInputStream(fileInputStream);
            fileOutputStream = new FileOutputStream(desUnzipFilePath);
            //读入压缩文件对象
            byte[] bytes = (byte[]) objectInputStream.readObject();
            //读入Huffman编码对应对象
            Map<Byte, String> huffmanCode = (Map<Byte, String>) objectInputStream.readObject();
            //解压文件
            byte[] unzipFile = unZipHuffmanCode(huffmanCode, byteToBinaryString(bytes));
            fileOutputStream.write(unzipFile);
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            try {
                //关闭流对象
                assert fileOutputStream != null;
                fileOutputStream.close();
                objectInputStream.close();
                fileInputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        System.out.println("解压完成");
    }

    /**
     * 压缩文件
     *
     * @param sourceFilePath 源文件路径全路径
     * @param desFilePath    压缩文件的目的全路径
     */
    public static void fileHuffmanZip(String sourceFilePath, String desFilePath) {
        FileInputStream fileInputStream = null;
        FileOutputStream fileOutputStream = null;
        ObjectOutputStream objectOutputStream = null;

        try {
            fileInputStream = new FileInputStream(sourceFilePath);
            //创建字符数组存储待压缩文件
            int len = 0;
            byte[] bytes = new byte[fileInputStream.available()];
//            while ((len = fileInputStream.read(bytes)) != -1) {
//                fileInputStream.read(bytes, 0, len);
//            }
            fileInputStream.read(bytes);
            //压缩文件
            byte[] zipHuffmanCode = huffmanZip(new String(bytes));
            fileOutputStream = new FileOutputStream(desFilePath);
            //创建对象输出流，存储文件
            objectOutputStream = new ObjectOutputStream(fileOutputStream);
            //存储编码后（压缩后）的Huffman编码
            objectOutputStream.writeObject(zipHuffmanCode);
            //存储编码对应关系
            objectOutputStream.writeObject(huffmanCodeMap);
            System.out.println("压缩完成！");
            System.out.println(byteToBinaryString(zipHuffmanCode));
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            //关闭输入输出流
            try {
                assert objectOutputStream != null;
                objectOutputStream.close();
                fileOutputStream.close();
                fileInputStream.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 解压赫夫曼编码
     *
     * @param huffmanCodeMap 赫夫曼编码后的字符与编码对应map
     * @param str            待解压的二进制字符串
     * @return 返回解压后的字符数组
     */
    public static byte[] unZipHuffmanCode(Map<Byte, String> huffmanCodeMap, String str) {
        HashMap<String, Byte> reverseHuffmanCodeMap = new HashMap<>();
        //反转huffmanCodeMap
        huffmanCodeMap.forEach((key, value) -> {
            reverseHuffmanCodeMap.put(value, key);
        });
        //存放匹配到字符
        ArrayList<Byte> bytes = new ArrayList<>();
        //匹配二进制字符串在反转后map中的value
        for (int i = 0; i < str.length(); ) {
            int count = 0;//索引，移动遍历
            boolean flag = true;//标志位，标志是否匹配到
            while (flag) {
                Byte huffmanCode = reverseHuffmanCodeMap.get(str.substring(i, i + count));
                if (huffmanCode == null) {
                    //未匹配到
                    count++;
                } else {
                    //匹配到
                    bytes.add(huffmanCode);
                    //退出循环
                    flag = false;
                }
            }
            //i指针移动到count位置
            i += count;
        }
        //将集合内的字符串重新放入byte数组
        byte[] huffmanBytes = new byte[bytes.size()];
        for (int i = 0; i < bytes.size(); i++) {
            huffmanBytes[i] = bytes.get(i);
        }
        return huffmanBytes;
    }

    /**
     * 将数组中的字符转为二进制字符串
     *
     * @param bytes 待转换数组
     * @return 二进制字符串组合
     */
    public static String byteToBinaryString(byte[] bytes) {
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < bytes.length; i++) {
            int temp = bytes[i];
            if (i != bytes.length - 1) {
                //对不是最后一位的字符（主要是正数）进行补位 256=》1 0000 0000
                temp |= 256;
            }
            //转为二进制字符串
            String binaryStr = Integer.toBinaryString(temp);
            if (i != bytes.length - 1 || temp < 0) {
                //不是最后一位,截取后八位
                stringBuilder.append(binaryStr.substring(binaryStr.length() - 8));
            } else {
//                if (temp > 1) {
//                    int i1 = Integer.parseInt(binaryStr, 2);
//                    i1 |= 256;
//                    String s = Integer.toBinaryString(i1);
//                    stringBuilder.append(s.substring((s.length() - binaryStr.length()) - 1));
//                } else {
//                }
                //TODO 在末尾是正数时会存在问题，未修改
                stringBuilder.append(binaryStr);
                //是最后一位，什么也不做
            }
        }
        return stringBuilder.toString();
    }

    /**
     * 重新封装Huffman编码压缩
     *
     * @param str 待压缩的字符串
     * @return 压缩后的字符数组
     */
    public static byte[] huffmanZip(String str) {
        byte[] bytes = str.getBytes();
        List<HuffmanTreeCodeNode> huffmanTreeCodeNodes = countStr(bytes);
        HuffmanTreeCodeNode huffmanTree = createHuffmanTree(huffmanTreeCodeNodes);
        Map<Byte, String> huffmanCode = createHuffmanCode(huffmanTree);
        return HuffmanCodeZip(huffmanCode, bytes);
    }

    /**
     * 统计字符出现的次数，并创建节点对象
     *
     * @param str 字符数组
     * @return 节点对象集合
     */
    public static List<HuffmanTreeCodeNode> countStr(byte[] str) {
        //定义map存储字符出现的次数
        HashMap<Byte, Integer> huffmanHashMaps = new HashMap<>();
        ArrayList<HuffmanTreeCodeNode> huffmanTreeCodeNodes = new ArrayList<>();
        //遍历byte数组
        for (byte node : str) {
            Integer count = huffmanHashMaps.get(node);
            if (count == null) {
                //第一次统计，赋值为1
                huffmanHashMaps.put(node, 1);
            } else {
                huffmanHashMaps.put(node, ++count);
            }
        }
        //遍历hashMap
        huffmanHashMaps.forEach((key1, value) -> {
            //放入集合内
            huffmanTreeCodeNodes.add(new HuffmanTreeCodeNode(value, key1));
        });
        return huffmanTreeCodeNodes;
    }

    /**
     * 创建赫夫曼树
     *
     * @param huffmanTreeCodeNodes 赫夫曼节点集合
     */
    public static HuffmanTreeCodeNode createHuffmanTree(List<HuffmanTreeCodeNode> huffmanTreeCodeNodes) {
        //遍历节点集合直到集合内只剩下根节点为止
        while (huffmanTreeCodeNodes.size() > 1) {
            //排序集合
            Collections.sort(huffmanTreeCodeNodes);
            //取出最小的两个，并删除
            HuffmanTreeCodeNode leftNode = huffmanTreeCodeNodes.remove(0);
            HuffmanTreeCodeNode rightNode = huffmanTreeCodeNodes.remove(0);
            //组成新的节点
            HuffmanTreeCodeNode huffmanTreeCodeNode = new HuffmanTreeCodeNode(leftNode.weight + rightNode.weight, null);
            //建立关系
            huffmanTreeCodeNode.leftNode = leftNode;
            huffmanTreeCodeNode.rightNode = rightNode;
            //放入集合内
            huffmanTreeCodeNodes.add(huffmanTreeCodeNode);
        }
        return huffmanTreeCodeNodes.get(0);
    }

    /**
     * 前序遍历赫夫曼树
     */
    public static void preOrder(HuffmanTreeCodeNode node) {
        if (node == null) {
            System.out.println("树为空");
            return;
        }
        node.preOrder();
    }

    //存储赫夫曼编码和data
    public static Map<Byte, String> huffmanCodeMap = new HashMap<>();

    /**
     * 方法重载
     *
     * @param node 根节点
     * @return 存储了赫夫曼编码的集合
     */
    public static Map<Byte, String> createHuffmanCode(HuffmanTreeCodeNode node) {
        createHuffmanCode(node.leftNode, "0", new StringBuilder());
        createHuffmanCode(node.rightNode, "1", new StringBuilder());
        return huffmanCodeMap;
    }

    /**
     * 创建赫夫曼树编码方法
     *
     * @param node          待生成编码节点
     * @param code          编码
     * @param stringBuilder 上一步添加了编码的stringBuilder
     */
    public static void createHuffmanCode(HuffmanTreeCodeNode node, String code, StringBuilder stringBuilder) {
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        stringBuilder1.append(code);
        if (node != null) {
            if (node.data == null) {
                //非叶子节点
                //递归左子树
                createHuffmanCode(node.leftNode, "0", stringBuilder1);
                //递归右子树
                createHuffmanCode(node.rightNode, "1", stringBuilder1);
            } else {
                //叶子节点
                huffmanCodeMap.put(node.data, stringBuilder1.toString());
            }
        }
    }

    /**
     * 使用Huffman编码压缩
     *
     * @param huffmanCode 生成的编码集合
     * @param str         待编码的字符数组
     * @return 压缩后的字符编码
     */
    public static byte[] HuffmanCodeZip(Map<Byte, String> huffmanCode, byte[] str) {
        StringBuilder stringBuilder = new StringBuilder();
        for (byte ch : str) {
            //遍历生成的赫夫曼编码，放入stringBuilder
            stringBuilder.append(huffmanCode.get(ch));
        }
        String huffmanCodeStr = stringBuilder.toString();
        System.out.println(huffmanCodeStr);
        //创建存储新code的数组
        int len;
        if (huffmanCodeStr.length() % 8 == 0) {
            len = huffmanCodeStr.length() / 8;
        } else {
            len = huffmanCodeStr.length() / 8 + 1;
        }
        byte[] huffmanBytes = new byte[len];
        int index = 0;
        for (int i = 0; i < huffmanCodeStr.length(); i = i + 8) {
            //八位一次，转为二进制重新存入新的byte数组
            String substring;
            if (i + 8 < huffmanCodeStr.length()) {
                substring = huffmanCodeStr.substring(i, i + 8);
            } else {
                substring = huffmanCodeStr.substring(i);
            }
            //将二进制转为十进制，并重新放入新的byte数组
            huffmanBytes[index] = (byte) Integer.parseInt(substring, 2);
            index++;
        }
        return huffmanBytes;
    }
}

class HuffmanTreeCodeNode implements Comparable<HuffmanTreeCodeNode> {
    public int weight;//节点权重，字符出现的次数
    public Byte data;//节点数据，一个字符
    public HuffmanTreeCodeNode leftNode;//左节点
    public HuffmanTreeCodeNode rightNode;//右节点

    public HuffmanTreeCodeNode(int weight, Byte data) {
        this.weight = weight;
        this.data = data;
    }

    @Override
    public String toString() {
        return "HuffmanTreeCodeNode{" +
                "weight=" + weight +
                ", data='" + data + '\'' +
                '}';
    }

    /**
     * 前序遍历赫夫曼树
     */
    public void preOrder() {
        System.out.println(this);
        if (this.leftNode != null) {
            this.leftNode.preOrder();
        }
        if (this.rightNode != null) {
            this.rightNode.preOrder();
        }
    }

    @Override
    public int compareTo(HuffmanTreeCodeNode o) {
        //从小到大
        return this.weight - o.weight;
    }
}
